"""
Currently this module can do:
 - read an arbitrary image file and find if it's a valid ISO9660 format, even if it's in RAW format:

    mode, frame_size = get_image_spec(filename)
    if mode is ISO9660:
        print "Is ISO9660"
    elif mode is YELLOW_MODE_1:
        print "Is yellow book, mode 1"
    elif mode is YELLOW_MODE_2:
        print "Is yellow book, mode 2"
    elif mode is GREEN_MODE_2_FORM_1:
        print "Is green book, mode 2, form 1"
    elif mode is GREEN_MODE_2_FORM_2:
        print "Is green book, mode 2, form 2"


 - can generate a CUE file if it can:
    
    try:
        target = open("test.cue", "w")
        cue = generate_cue(filename)
        target.write(cue)
        target.close()
        
    except CDImageFormatError:
        print >> sys.stderr, "Invalid format", filename
        
 
 - can convert a valid image to an ISO9660 format:
 
    fd = extract_iso9660(filename)
    iso9660 = fd.read()
    fd.close()
 
 
"""
__author__ = "Tiago Cogumbreiro <cogumbreiro@users.sf.net>"
__copyright__ = "Copyright 2006, Tiago Cogumbreiro"
__license__ = "MIT <http://www.opensource.org/licenses/mit-license.php>"

from StringIO import StringIO

class CDImageSizeError(StandardError):
    """This error is raised when the image file is of a given file is not divided in chunks."""

class CDImageFormatError(StandardError):
    """This error is raised when the given file is not a CDImage"""

class CDImageModeError(StandardError):
    """This error is thrown when the given mode is invalid"""
    
# Constants

# Image Modes:
YELLOW_MODE_1 = "yellow_mode_1"
YELLOW_MODE_2 = "yellow_mode_2"
GREEN_MODE_2_FORM_1 = "green_mode_2_form_1"
GREEN_MODE_2_FORM_2 = "green_mode_2_form_2"
GREEN_MODE_2_FORM_2_SPARE = "green_mode_2_form_2_spare"
ISO9660 = "iso9660"



def is_iso9660(buff):
    """Checks if a given buffer is an ISO, by doing that it 32776 bytes"""
    try:
        white = buff.read(0x8000)
        if len(white) != 0x8000:
            return False

        cd_001 = "\x01CD001\x01\x00"
    
        return buff.read(8) == cd_001
    except CDImageSizeError:
        return False

def safe_read_userdata(buff, bytes_in, bytes_out, offset):
    chunk = buff.read(bytes_in)
    if len(chunk) != bytes_in:
        raise CDImageSizeError("Expecting %d but read %d bytes." % (bytes_in, len(chunk)))
    return chunk[offset:][:bytes_out]

def unsafe_read_userdata(buff, bytes_in, bytes_out, offset):
    return buff.read(bytes_in)[offset:][:bytes_out]

class CDImageReader:
    """
    Extracts userdata from a certain image.
    """
    
    MODES = {
        ISO9660: dict(offset=0, size=2048),
        YELLOW_MODE_1: dict(offset=16, size=2048),
        YELLOW_MODE_2: dict(offset=16, size=2336),
        GREEN_MODE_2_FORM_1: dict(offset=24, size=2048),
        GREEN_MODE_2_FORM_2: dict(offset=24, size=2324),
        GREEN_MODE_2_FORM_2_SPARE: dict(offset=24, size=2328),
    }

    def __init__(self, buff, mode, frame_size, check_size=True):
        try:
            self.mode = self.MODES[mode]
        except KeyError:
            raise CDImageModeError(mode)
            
        self.buff = buff
        self.frame_size = frame_size
        self.remaining_data = ""
        
        if check_size:
            self.read_userdata = safe_read_userdata
        else:
            self.read_userdata = unsafe_read_userdata
    
    def read(self, size=None):
        if size is None:
            # Read until finish
            data = buff.read()
            data_size = len(data)
            blocks = data_size / self.frame_size
            bytes_to_read = blocks * self.mode["size"]
            missing_bytes = data_size % frame_size
            offset = self.mode["offset"]
            frame_size = self.mode["frame_size"]
            
            if missing_bytes > offset + frame_size:
                missing_bytes = offset + frame_size
                
            if missing_bytes > self.mode["offset"]:
                bytes_to_read -= self.mode["offset"]
            return CDImageReader(StringIO(data), self.mode, self.frame_size).read(bytes_to_read)
        
        if len(self.remaining_data) >= size:
            data = self.remaining_data[:size]
            self.remaining_data = self.remaining_data[size:]
            return data
        
        data = self.remaining_data
        size -= len(data)
        
        data_size = self.mode["size"]
        offset = self.mode["offset"]
        
        full_blocks = size / data_size
        frame_size = self.frame_size
        
        # Read the full blocks
        read_data = lambda x=1: self.read_userdata(self.buff, frame_size, data_size, offset)
        data += "".join(map(read_data, range(full_blocks)))

        # Now read last chunk's bytes
        remaining_bytes = size % data_size
        if remaining_bytes > 0:
            self.remaining_data = read_data()
            return data + self.read(size % data_size)
        else:
            return data

    def close(self):
        self.buff.close()

def check_mode(filename, mode, frame_size=2352):
    """
    Checks if a given filename is of target mode and frame_size.
    """
    try:
        fd = open(filename)
    except IOError:
        return False

    try:
        if mode is ISO9660:
            assert frame_size == 2048
            filter = fd
        else:
            filter = CDImageReader(fd, mode, frame_size)

        try:
            if is_iso9660(filter):
                fd.close()
                return True
        except CDImageSizeError:
            pass
            
    finally:
        fd.close()
        
    return False



def get_image_spec(filename):
    """
    Returns the image mode and the frame size in a tuple.
    (mode, frame_size)

    Throws CDImageFormatError if it's an invalid CDImageFormat.
    """

    modes = (
        (ISO9660, 2048),
        (YELLOW_MODE_1, 2352),
        (YELLOW_MODE_2, 2352),
        (GREEN_MODE_2_FORM_1, 2352),
        (GREEN_MODE_2_FORM_2, 2352),
        (YELLOW_MODE_1, 2336),
    )
    
    for mode, frame_size in modes:
        if check_mode(filename, mode, frame_size):
            return mode, frame_size
    
    raise CDImageFormatError




def generate_cue(filename):
    """
    Generates a CUE from a given filename.
    Raises CDImageFormatError if it cannot generate it.

    A CUE/BIN file supports 3 types of modes:
    Name       - Type, frame_size
    MODE1/2048 - ISO9660             - 2048
    MODE1/2352 - YELLOW_MODE_1       - 2352
    MODE2/2352 - GREEN_MODE_2_FORM_1 - 2352
    MODE2/2336 - YELLOW_MODE_1       - 2336
    """
    # These are the mode/frame compatible with the CUE/BIN format
    cue_header = {
        (ISO9660, 2048): "MODE1/2048",
        (YELLOW_MODE_1, 2352): "MODE1/2352",
        (GREEN_MODE_2_FORM_1, 2352): "MODE2/2352",
        (YELLOW_MODE_1, 2336): "MODE2/2336",
    }
    
    try:
        mode = cue_header[get_image_spec(filename)]
    except KeyError:
        raise CDImageFormatError
    
    cue = """FILE "%(filename)s" BINARY
TRACK 01 %(mode)s
INDEX 01 00:00:00
"""
    return cue % locals()




def extract_iso9660(filename, fail_if_iso9660=True):
    """If the filename is a CD Image it converts it to a ISO9660 image."""
    mode, frame_size = check_image(filename)
    fd = open(filename)
    
    if mode is ISO9660:
        if fail_if_iso9660:
            raise CDImageFormatError("File is already at ISO9660 format")
            
        return fd
    else:
        return CDImageReader(fd, mode, frame_size)
    



if __name__ == '__main__':
    import sys
    for filename in sys.argv[1:]:
        try:
            generate_cue(filename)
            print filename
        except CDImageFormatError:
            print >> sys.stderr, "Cannot generate a CUE from %s" % filename
        
#class TestCDImageReader:
